"""
基于机器学习的提示词质量评估模块
使用sentence-transformers模型辅助评估提示词质量
"""

import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from sklearn.metrics.pairwise import cosine_similarity


class MLPromptEvaluator:
    """
    基于机器学习的提示词评估器
    使用嵌入模型计算语义相似度来评估提示词质量
    """
    
    def __init__(self):
        self._embedding_model = None
        self._high_quality_prompts = []
        self._initialize_model()
    
    def _initialize_model(self):
        """初始化嵌入模型"""
        try:
            from sentence_transformers import SentenceTransformer
            self._embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        except ImportError:
            print("警告: sentence_transformers 未安装，将使用简单的文本匹配")
            self._embedding_model = None
    
    def add_high_quality_examples(self, prompts: List[str]):
        """
        添加高质量提示词示例用于参考比较
        
        Args:
            prompts: 高质量提示词列表
        """
        self._high_quality_prompts.extend(prompts)
        if self._embedding_model and prompts:
            # 预计算高质量提示词的嵌入
            embeddings = self._embedding_model.encode(prompts)
            self._high_quality_embeddings = embeddings
    
    def evaluate_prompt(self, prompt: str, requirement: str = "") -> Dict[str, Any]:
        """
        使用机器学习模型评估提示词质量
        
        Args:
            prompt: 要评估的提示词
            requirement: 原始需求（可选）
            
        Returns:
            包含评估结果的字典
        """
        if not prompt:
            return {
                "ml_score": 0.0,
                "similarity_to_examples": 0.0,
                "semantic_coherence": 0.0,
                "details": "提示词为空"
            }
        
        if not self._embedding_model:
            # 如果没有嵌入模型，使用简单的启发式方法
            return self._simple_evaluation(prompt, requirement)
        
        try:
            # 计算提示词嵌入
            prompt_embedding = self._embedding_model.encode([prompt])
            
            # 计算与高质量示例的相似度
            similarity_to_examples = 0.0
            if hasattr(self, '_high_quality_embeddings') and len(self._high_quality_embeddings) > 0:
                similarities = cosine_similarity(
                    prompt_embedding, 
                    self._high_quality_embeddings
                )
                similarity_to_examples = float(np.mean(similarities))
            
            # 计算语义连贯性（如果提供了需求）
            semantic_coherence = 0.0
            if requirement:
                req_embedding = self._embedding_model.encode([requirement])
                coherence_sim = cosine_similarity(prompt_embedding, req_embedding)
                semantic_coherence = float(coherence_sim[0][0])
            
            # 综合得分
            ml_score = (similarity_to_examples * 0.6 + semantic_coherence * 0.4)
            
            return {
                "ml_score": round(ml_score, 3),
                "similarity_to_examples": round(similarity_to_examples, 3),
                "semantic_coherence": round(semantic_coherence, 3),
                "details": "基于语义嵌入的评估完成"
            }
            
        except Exception as e:
            return {
                "ml_score": 0.0,
                "similarity_to_examples": 0.0,
                "semantic_coherence": 0.0,
                "details": f"评估失败: {str(e)}"
            }
    
    def _simple_evaluation(self, prompt: str, requirement: str = "") -> Dict[str, Any]:
        """
        简单的启发式评估方法（当无法使用嵌入模型时）
        
        Args:
            prompt: 要评估的提示词
            requirement: 原始需求（可选）
            
        Returns:
            包含评估结果的字典
        """
        # 基于长度的简单评估
        length_score = min(len(prompt) / 500, 1.0)
        
        # 基于关键词的简单评估
        quality_keywords = ["明确", "具体", "要求", "输出", "示例", "格式", "步骤"]
        keyword_score = sum(1 for kw in quality_keywords if kw in prompt) / len(quality_keywords)
        
        # 基于结构的简单评估
        structure_indicators = ["1.", "2.", "第一", "第二", "- ", "* "]
        structure_score = min(1.0, sum(1 for si in structure_indicators if si in prompt) / 3)
        
        # 如果提供了需求，进行简单的文本匹配
        requirement_match = 0.0
        if requirement:
            # 简单的词重叠计算
            prompt_words = set(prompt.split())
            req_words = set(requirement.split())
            if req_words:
                requirement_match = len(prompt_words.intersection(req_words)) / len(req_words)
        
        # 综合得分
        ml_score = (length_score * 0.3 + keyword_score * 0.3 + structure_score * 0.2 + requirement_match * 0.2)
        
        return {
            "ml_score": round(ml_score, 3),
            "similarity_to_examples": round(keyword_score, 3),
            "semantic_coherence": round(requirement_match, 3),
            "details": "基于启发式的简单评估"
        }
    
    def compare_prompts(self, prompts: List[str]) -> Dict[str, Any]:
        """
        比较多个提示词的质量
        
        Args:
            prompts: 提示词列表
            
        Returns:
            包含比较结果的字典
        """
        if not prompts:
            return {"error": "没有提供提示词"}
        
        if not self._embedding_model:
            return {"error": "嵌入模型不可用"}
        
        try:
            # 计算所有提示词的嵌入
            embeddings = self._embedding_model.encode(prompts)
            
            # 计算两两之间的相似度
            similarities = cosine_similarity(embeddings)
            
            # 计算每个提示词的平均相似度（与其他提示词的相似度）
            avg_similarities = np.mean(similarities, axis=1)
            
            # 找到最佳提示词（平均相似度最高的）
            best_index = int(np.argmax(avg_similarities))
            
            return {
                "similarities": similarities.tolist(),
                "avg_similarities": avg_similarities.tolist(),
                "best_prompt_index": best_index,
                "best_prompt": prompts[best_index]
            }
            
        except Exception as e:
            return {"error": f"比较失败: {str(e)}"}


def create_ml_evaluator_with_examples() -> MLPromptEvaluator:
    """
    创建带有高质量示例的ML评估器
    
    Returns:
        配置好的MLPromptEvaluator实例
    """
    evaluator = MLPromptEvaluator()
    
    # 添加一些高质量提示词示例
    high_quality_examples = [
        """你是一位经验丰富的Python开发专家。请根据以下要求编写代码：
        1. 功能需求：实现一个函数，用于计算斐波那契数列的前n项
        2. 技术要求：
           - 使用递归算法实现
           - 添加输入验证，确保n为非负整数
           - 包含异常处理
           - 添加详细的中文注释说明算法逻辑
        3. 输出要求：
           - 提供完整的可执行代码
           - 代码遵循PEP8规范
           - 包含简单的测试用例
        4. 示例：
           输入：n=5
           输出：[0, 1, 1, 2, 3, 5]""",
        
        """作为数据分析师，请按照以下步骤分析给定的数据集：
        1. 数据概览：
           - 描述数据集的基本信息（行数、列数、数据类型等）
           - 检查缺失值情况
        2. 数据清洗：
           - 处理缺失值
           - 处理异常值
        3. 数据分析：
           - 计算基本统计量
           - 分析变量间的关系
        4. 数据可视化：
           - 创建适当的图表展示关键发现
        5. 结论总结：
           - 总结主要发现
           - 提出建议""",
        
        """你是一位专业的技术文档撰写专家。请为以下技术概念编写清晰易懂的解释文档：
        1. 概念定义：准确、简洁地定义该技术概念
        2. 核心原理：解释其工作原理和核心机制
        3. 应用场景：说明在哪些情况下使用该技术
        4. 优缺点分析：列出主要优点和局限性
        5. 实际示例：提供具体的使用示例
        6. 相关概念：介绍相关的技术概念及其关系
        7. 学习资源：推荐进一步学习的资料"""
    ]
    
    evaluator.add_high_quality_examples(high_quality_examples)
    return evaluator


# 导出类和函数
__all__ = [
    'MLPromptEvaluator',
    'create_ml_evaluator_with_examples'
]